from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
import httpx
import os
from pydantic import BaseModel
from typing import List, Optional
import json
import logging

router = APIRouter(prefix="/gpt")

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 定义请求和响应模型
class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    model: Optional[str] = "deepseek-chat"
    temperature: Optional[float] = 0.7

# DeepSeek API配置
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "sk-ce4c94f281b74894bb4a3fcf0a46ef9d")
DEEPSEEK_API_URL = "https://api.deepseek.com/chat/completions"

@router.post("/polish_txt_stream")
async def polish_text_stream(request: ChatRequest):
    """
    调用DeepSeek Chat接口实现流式对话功能 (SSE格式)
    """
    try:
        headers = {
            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "model": request.model,
            "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
            "temperature": request.temperature,
            "stream": True
        }
        
        async def generate():
            try:
                async with httpx.AsyncClient() as client:
                    async with client.stream(
                        "POST",
                        DEEPSEEK_API_URL,
                        headers=headers,
                        json=payload,
                        timeout=30.0
                    ) as response:
                        if response.status_code != 200:
                            error_msg = await response.aread()
                            logger.error(f"DeepSeek API error: {error_msg.decode()}")
                            error_data = {
                                "error": {
                                    "message": f"DeepSeek API error: {response.status_code}",
                                    "type": "api_error",
                                    "code": response.status_code
                                }
                            }
                            yield f"data: {json.dumps(error_data)}\n\n"
                            yield "data: [DONE]\n\n"
                            return
                        
                        # sse 流式传输
                        # 开始：data: 
                        # 接口：[DONE]
                        
                        async for line in response.aiter_lines():
                            if line.startswith("data: "):
                                data = line[6:]
                                if data.strip() == "[DONE]":
                                    yield "data: [DONE]\n\n"
                                    break
                                # 验证是否为有效的JSON
                                try:
                                    json.loads(data)
                                    yield f"data: {data}\n\n"
                                except json.JSONDecodeError:
                                    # 跳过无效的JSON行
                                    continue
                                    
            except Exception as e:
                logger.error(f"Internal server error: {str(e)}")
                error_data = {
                    "error": {
                        "message": f"Internal server error: {str(e)}",
                        "type": "internal_error",
                        "code": 500
                    }
                }
                yield f"data: {json.dumps(error_data)}\n\n"
                yield "data: [DONE]\n\n"
        
        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "Content-Type": "text/event-stream; charset=utf-8",
                "Access-Control-Allow-Origin": "*",
            }
        )
            
    except Exception as e:
        logger.error(f"Failed to process request: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Internal server error: {str(e)}"
        )